import torch
import numpy as np

from data.transform import TransformLoader


class MapDataset(torch.utils.data.Dataset):
    def __init__(self, dataset, mapper):
        self.dataset = dataset
        self.mapper = mapper

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        data = self.dataset[index]
        data = self.mapper(data)
        return data

    def __iter__(self):
        for i in range(len(self)):
            yield self[i]


class DataMapper:
    def __init__(self, loader_config):
        self.size = loader_config['input_size']
        self.categories = loader_config['categories']
        self.addition_transforms = loader_config['transforms'] + [
            {
                "type": "MapCategoryToId",
                "categories": self.categories,
            },
            {
                "type": "ResizePadding",
                "size": self.size,
            },
        ]

    def __call__(self, data):
        data['transforms'] += self.addition_transforms

        transform = TransformLoader.get(*data['transforms'])
        data = transform(data)

        image = data['image']
        if len(image.shape) == 2:
            image = np.expand_dims(image, -1)
        assert len(image.shape) == 3

        data['image'] = image.transpose(2, 0, 1)

        return data


class MultiScaleDataMapper:
    def __init__(self, loader_config):
        self.sizes = loader_config['input']['sizes'] if 'sizes' in loader_config['input'] else [loader_config['input']['size']]
        self.size = None
        self.batch_size = loader_config['batch_size']
        self.count = 0

        self.categories = loader_config['categories']
        self.addition_transforms = loader_config['transforms'] + [
            {
                "type": "MapCategoryToId",
                "categories": self.categories,
            },
        ]

    def __call__(self, data):
        if self.count % self.batch_size == 0:
            self.size = self.sizes[np.random.choice(len(self.sizes))]
            self.count = 0
        self.count += 1

        data['transforms'] += self.addition_transforms + [
            {
                "type": "ResizePadding",
                "size": self.size,
            }
        ]

        transform = TransformLoader.get(*data['transforms'])
        data = transform(data)

        image = data['image']
        if len(image.shape) == 2:
            image = np.expand_dims(image, -1)
        assert len(image.shape) == 3

        data['image'] = image.transpose(2, 0, 1)

        return data
